package com.zhxq.ai.websocket;

import jakarta.websocket.OnMessage;
import jakarta.websocket.OnOpen;
import jakarta.websocket.Session;
import jakarta.websocket.server.PathParam;
import jakarta.websocket.server.ServerEndpoint;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

@Slf4j
@Component
@EnableWebSocket
@ServerEndpoint("/member/ws/interview/{chatId}")
@RequiredArgsConstructor
public class WebSocketInterviewServer {

    private static Map<String, Session> sessionMap = new HashMap<>();

    @OnOpen
    public void onOpen(Session session, @PathParam("chatId") String chatId) {
        log.info("WebSocketServer onOpen {}", session.getId());
        sessionMap.put(chatId, session);
        log.info("chatId: {}", chatId);
    }

    @OnMessage
    public void onMessage(Session session, String message) {
        log.info("接收到了客户端 {} 发来的消息 : {}", session.getId(), message);
    }

    public void sendMessageToAI(String chatId, Flux<String> responseStream) throws InterruptedException {
        log.info("WebSocketServer sendMessageToAI {}", chatId);
        Session session = sessionMap.get(chatId);
        responseStream.subscribe(
                data -> {
                    try {
                        session.getBasicRemote().sendText(data);
                    } catch (Exception e) {
                        throw new RuntimeException(e);
                    }
                },
                err -> {
                    try {
                        session.getBasicRemote().sendText("[ERROR]" + err.getMessage());
                    } catch (Exception e) {
                        throw new RuntimeException(e);
                    }
                },
                () -> {
                    try {
                        session.getBasicRemote().sendText("[END]"); // 发送结束标识
                    } catch (IOException e) {
                        throw new RuntimeException(e);
                    }
                }
        );
    }
}
